import dotenv
import invoke
from langchain.chains.conversation.base import ConversationChain
from langchain.memory import ConversationEntityMemory
from langchain.memory.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE
from langchain_community.chat_models import ChatOpenAI

dotenv.load_dotenv()
# 构建提示词
llm = ChatOpenAI(model="gpt-4o-mini",temperature=0)
# llm = ChatOpenAI(model="gpt-3.5-turbo-16k",temperature=0)


chain = ConversationChain(
    llm = llm,
    prompt = ENTITY_MEMORY_CONVERSATION_TEMPLATE,
    memory = ConversationEntityMemory(llm=llm)
)

print(chain.invoke({"input":"你好我叫博小瑞。我最近在学习langchain"}))
print(chain.invoke({"input":"我最喜欢的编程语言是Python"}))
print(chain.invoke({"input":"我住在广州"}))

res = chain.memory.entity_store.store

print(res)

# chain = ConversationChain